﻿using LDASharp;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace GibbsLDASharp.LDA
{
    public class Inferencer
    {
        // Train model
        public Model trnModel;
        public LDADictionary globalDict;
        private LDACmdOption option;

        private Model newModel;
        public int niters = 100;

        //-----------------------------------------------------
        #region Init method
        //-----------------------------------------------------
        public Boolean Init(LDACmdOption option)
        {
            this.option = option;
            trnModel = new Model();

            if (!trnModel.InitEstimatedModel(option))
                return false;

            globalDict = trnModel.data.localDict;
            ComputeTrnTheta();
            ComputeTrnPhi();

            return true;
        }
        #endregion
        //inference new model ~ getting data from a specified dataset
        public Model Inference(LDADataset newData)
        {
            Console.WriteLine("init new model");
            Model newModel = new Model();

            newModel.InitNewModel(option, newData, trnModel);
            this.newModel = newModel;

            Console.WriteLine("Sampling " + niters + " iteration for inference!");
            for (newModel.liter = 1; newModel.liter <= niters; newModel.liter++)
            {
                //Console.WriteLine("Iteration " + newModel.liter + " ...");

                // for all newz_i
                for (int m = 0; m < newModel.M; ++m)
                {
                    for (int n = 0; n < newModel.data.docs[m].length; n++)
                    {
                        // (newz_i = newz[m][n]
                        // sample from p(z_i|z_-1,w)
                        int topic = InfSampling(m, n);
                        newModel.z[m][n] = topic;
                    }
                }//end foreach new doc

            }// end iterations

            Console.WriteLine("Gibbs sampling for inference completed!");

            ComputeNewTheta();
            ComputeNewPhi();
            newModel.liter--;
            return this.newModel;
        }

        public Model Inference(String[] strs)
        {
            //Console.WriteLine("inference");
            Model newModel = new Model();

            //Console.WriteLine("read dataset");
            LDADataset dataset = LDADataset.ReadDataSet(strs, globalDict);

            return Inference(dataset);
        }

        //inference new model ~ getting dataset from file specified in option
        public Model Inference()
        {
            //Console.WriteLine("inference");

            newModel = new Model();
            if (!newModel.InitNewModel(option, trnModel)) return null;

            Console.WriteLine("Sampling " + niters + " iteration for inference!");

            for (newModel.liter = 1; newModel.liter <= niters; newModel.liter++)
            {
                //Console.WriteLine("Iteration " + newModel.liter + " ...");

                // for all newz_i
                for (int m = 0; m < newModel.M; ++m)
                {
                    for (int n = 0; n < newModel.data.docs[m].length; n++)
                    {
                        // (newz_i = newz[m][n]
                        // sample from p(z_i|z_-1,w)
                        int topic = InfSampling(m, n);
                        newModel.z[m][n] = topic;
                    }
                }//end foreach new doc

            }// end iterations

            Console.WriteLine("Gibbs sampling for inference completed!");
            Console.WriteLine("Saving the inference outputs!");

            ComputeNewTheta();
            ComputeNewPhi();
            newModel.liter--;
            newModel.SaveModel(newModel.dfile + "." + newModel.modelName);

            return newModel;
        }

        



        /// <summary>
        /// Do sampling for inference
        /// </summary>
        /// <param name="m">Document number</param>
        /// <param name="n">Word number</param>
        /// <returns></returns>
        protected int InfSampling(int m, int n)
        {
            // remove z_i from the count variables
            int topic = newModel.z[m][n];
            int _w = newModel.data.docs[m].words[n];
            int w = newModel.data.lid2gid[_w];
            newModel.nw[_w, topic] -= 1;
            newModel.nd[m, topic] -= 1;
            newModel.nwsum[topic] -= 1;
            newModel.ndsum[m] -= 1;

            double Vbeta = trnModel.V * newModel.beta;
            double Kalpha = trnModel.K * newModel.alpha;

            Random rand = new Random();

            // do multinomial sampling via cummulative method		
            for (int k = 0; k < newModel.K; k++)
            {
                newModel.p[k] = (trnModel.nw[w, k] + newModel.nw[_w, k] + newModel.beta) / (trnModel.nwsum[k] + newModel.nwsum[k] + Vbeta) *
                        (newModel.nd[m, k] + newModel.alpha) / (newModel.ndsum[m] + Kalpha);
            }

            // cummulate multinomial parameters
            for (int k = 1; k < newModel.K; k++)
            {
                newModel.p[k] += newModel.p[k - 1];
            }

            // scaled sample because of unnormalized p[]
            double u = rand.NextDouble() * newModel.p[newModel.K - 1];

            for (topic = 0; topic < newModel.K; topic++)
            {
                if (newModel.p[topic] > u)
                    break;
            }

            // add newly estimated z_i to count variables
            newModel.nw[_w, topic] += 1;
            newModel.nd[m, topic] += 1;
            newModel.nwsum[topic] += 1;
            newModel.ndsum[m] += 1;

            return topic;
        }

        protected void ComputeNewTheta()
        {
            for (int m = 0; m < newModel.M; m++)
            {
                for (int k = 0; k < newModel.K; k++)
                {
                    newModel.theta[m, k] = (newModel.nd[m, k] + newModel.alpha) / (newModel.ndsum[m] + newModel.K * newModel.alpha);
                }//end foreach topic
            }//end foreach new document
        }

        protected void ComputeNewPhi()
        {
            for (int k = 0; k < newModel.K; k++)
            {
                for (int _w = 0; _w < newModel.V; _w++)
                {
                    int id;
                    if (newModel.data.lid2gid.TryGetValue(_w, out id))
                    {
                        newModel.phi[k, _w] = (trnModel.nw[id, k] + newModel.nw[_w, k] + newModel.beta) / (newModel.nwsum[k] + newModel.nwsum[k] + trnModel.V * newModel.beta);
                    }
                }//end foreach word
            }// end foreach topic
        }

        protected void ComputeTrnTheta()
        {
            for (int m = 0; m < trnModel.M; m++)
            {
                for (int k = 0; k < trnModel.K; k++)
                {
                    trnModel.theta[m, k] = (trnModel.nd[m, k] + trnModel.alpha) / (trnModel.ndsum[m] + trnModel.K * trnModel.alpha);
                }
            }
        }

        protected void ComputeTrnPhi()
        {
            for (int k = 0; k < trnModel.K; k++)
            {
                for (int w = 0; w < trnModel.V; w++)
                {
                    trnModel.phi[k, w] = (trnModel.nw[w, k] + trnModel.beta) / (trnModel.nwsum[k] + trnModel.V * trnModel.beta);
                }
            }
        }
    }
}
